/*
 * message_passing.h
 *
 *  Created on: Sep 12, 2011
 *      Author: tqlong
 */

#ifndef MESSAGE_PASSING_H_
#define MESSAGE_PASSING_H_

#include "simple_gm.h"

/** Classical message passing algorithm for
 *    + node & edge marginals
 *    + map assignment
 */
class MessagePassingGM : public SimpleGM
{
public:
  typedef std::map<Node, arma::vec>                    n_potential_type;
  typedef std::map<Edge, arma::mat, Edge::EdgeCompare> e_potential_type;
  typedef std::map<int, int>                           imap_type;
  typedef imap_type                                    assignment_type;
  typedef std::vector<Edge>                            evec_type;
  typedef std::map<Node, evec_type>                    adjacent_edge_type;

  typedef std::map<Edge, arma::vec, Edge::EdgeCompare> message_type;
  typedef const n_potential_type::value_type&          const_pair_node_t;
  typedef n_potential_type::value_type&                pair_node_t;
  typedef const e_potential_type::value_type&          const_pair_edge_t;
  typedef e_potential_type::value_type&                pair_edge_t;
  typedef const message_type::value_type&              const_pair_message_t;
  typedef message_type::value_type&                    pair_message_t;
public:
  /** Compute all node and edge marginals */
  void message_passing_marginals();
  /** Compute MAP assignment */
  assignment_type message_passing_map();

  MessagePassingGM reduce(const assignment_type& a)
  {
    return SimpleGM::reduce(a);
  }

  MessagePassingGM() {}
  MessagePassingGM(const SimpleGM& gm);

public:
  /** Message : m_ij(x_j) */
  message_type messages_;
  std::vector<Edge> schedule_;
  n_potential_type max_marginals_;

  /** Print information */
  void print_messages();
protected:
  void init_messages();
  void create_message(int i, int j);
  void create_passing_schedule();
  double message_diffence(const message_type& m1, const message_type& m2);

  void update_message(const Edge& e, message_type& msg);
  void message_to_node_marginals();
  void message_to_edge_marginals();

  void update_max_message(const Edge& e, message_type& msg);
  void message_to_max_marginals();
};

#endif /* MESSAGE_PASSING_H_ */
